import numpy as np
import matplotlib.pyplot as plt


from laplace_noise import add_laplace_noise, mean_confidence_interval

from dp.commons import TDP, PDP, RDP, set_fig_font

set_fig_font()

TRUTH = 10000


def plot(ax, xs, ys, errors, title, limit_y=True):
    ax.errorbar(xs, ys, yerr=errors, color="k", alpha=0.3, ls="None")
    ax.scatter(xs, ys, s=10 / np.log10(np.array(xs)), marker="o", color="r")
    ax.set_xscale("log")
    # ax.set_yscale("log")
    ax.set_xlim([0.95, 10010])
    if limit_y:
        ax.set_ylim([10000 - 30, 10000 + 30])
    ax.axhline(y=TRUTH, color="gray", linestyle="dotted")
    ax.text(1, TRUTH, r"\textbf{truth=10k}")
    ax.set_title(title, y=-0.3)


if __name__ == "__main__":

    sz = 10000

    xs = [
        *range(3, 10),
        *range(10, 30, 2),
        *range(30, 40, 3),
        *range(40, 50, 5),
        *range(50, 100, 8),
        *range(100, 300, 10),
        *range(300, 500, 20),
        *range(500, 1000, 30),
        *range(1000, 10000, 50),
    ]
    # xs= [*range(3,15)]
    fig, axs = plt.subplots(3, figsize=(7, 8))

    data = np.full(sz, TRUTH)
    noisy_data = add_laplace_noise(data, sensitivity=1, epsilon=0.1)
    dps = []
    errors = []
    for i in xs:
        dp_laplace, error_laplace = mean_confidence_interval(noisy_data, i)
        dps.append(dp_laplace)
        errors.append(error_laplace)
    plot(axs[0], xs, dps, errors, "(a) " + TDP)

    data = np.full(sz, TRUTH)
    noisy_data = add_laplace_noise(
        data, sensitivity=1, epsilon=0.1, dynamic_epsilon=True
    )
    dps = []
    errors = []
    for i in xs:
        dp_laplace, error_laplace = mean_confidence_interval(noisy_data, i)
        dps.append(dp_laplace)
        errors.append(error_laplace)
    plot(axs[1], xs, dps, errors, "(b) " + PDP)

    data = np.full(sz, TRUTH)
    noisy_data = add_laplace_noise(data, sensitivity=1, epsilon=0.1, composition=True)
    dps = []
    errors = []
    for i in xs:
        dp_laplace, error_laplace = mean_confidence_interval(noisy_data, i)
        dps.append(dp_laplace)
        errors.append(error_laplace)
    plot(axs[2], xs, dps, errors, "(c) " + "Composition DP", limit_y=False)

    # data = np.full(sz, truth)
    # noisy_data = add_laplace_noise(data, sensitivity=1, epsilon=0.1, seed=1)
    # dps = []
    # errors = []
    # for i in xs:
    #     dp_laplace, error_laplace = mean_confidence_interval(noisy_data, i)
    #     dps.append(dp_laplace)
    #     errors.append(error_laplace)
    # plot(axs[2], xs, dps, errors, "(c) " + RDP, limit_y=False)

    # plt.xlabel("Number of queries")
    fig.text(0.5, 0.02, "Number of queries", ha="center")
    fig.text(0.04, 0.5, "Estimation with 95\% CI", va="center", rotation="vertical")
    # plt.tight_layout()
    plt.subplots_adjust(hspace=0.3)
    plt.show()
